{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Examples\n",
    "System regression simultaneously estimates multiple models.  This has three distinct advantages:\n",
    "\n",
    "* Joint inference across models\n",
    "* Linear restrictions can be imposed on the parameters across different models\n",
    "* Improved precision of parameter estimates (depending on the model specification and data)\n",
    "\n",
    "There are $K$ models and each model can be expressed in vector notation as \n",
    "\n",
    "$$ Y_i = X_i\\beta_i + \\epsilon_i$$\n",
    "\n",
    "so that the set of models can be expressed as \n",
    "\n",
    "$$ Y = X\\beta + \\epsilon$$\n",
    "\n",
    "where $Y$ is a column vector that stacks the vectors $Y_i$ for $i=1,2,\\ldots,K$, $X$ is a block-diagonal matrix where diagonal block i is $X_i$, $\\beta$ is a stacked vector of the $K$ $\\beta_i$s and $\\epsilon$ is similarly comprised of the stacked columns of $\\epsilon_i$.\n",
    "\n",
    "The model can be estimated using OLS with the usual estimator\n",
    "\n",
    "$$\\hat{\\beta}_{OLS} = \\left(X^\\prime X\\right)^{-1}X^\\prime Y.$$\n",
    "\n",
    "Since there are multiple series, a GLS estimator that accounts for the cross-sectional heteroskedasticity as well as the correlation of residuals can be estimated \n",
    "\n",
    "$$\\hat{\\beta}_{GLS} = \\left(X^\\prime \\Omega^{-1} X\\right)^{-1}X^\\prime  \\Omega^{-1} Y$$\n",
    "\n",
    "where $\\Omega^{-1} = \\Sigma^{-1} \\otimes I_{T}$, $\\Sigma_{ij}$ is the covariance between $\\epsilon_i$ and $\\epsilon_j$ and $T$ is the number of observations. The GLS estimator is only beneficial when the regressors in different models differ  and when residuals are correlated. There GLS estimates are identical to the multivariate OLS estimates when all regressors are common."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Common libraries\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import statsmodels.api as sm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "Two data sets will be used.  The first is from Munnell which looks at the effect of capital on state GDP.  This example follows the example in Chapter 10 in recent editions of Greene\"s _Econometric Analysis_.\n",
    "\n",
    "The data is state-level but the model is estimated in region.  The first step is to aggregate the data by region.  All capital measures are summed and the unemployment rate is averaged using weights proportional to the total employment in each state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from linearmodels.datasets import munnell\n",
    "\n",
    "data = munnell.load()\n",
    "\n",
    "regions = {\n",
    "    \"GF\": [\"AL\", \"FL\", \"LA\", \"MS\"],\n",
    "    \"MW\": [\"IL\", \"IN\", \"KY\", \"MI\", \"MN\", \"OH\", \"WI\"],\n",
    "    \"MA\": [\"DE\", \"MD\", \"NJ\", \"NY\", \"PA\", \"VA\"],\n",
    "    \"MT\": [\"CO\", \"ID\", \"MT\", \"ND\", \"SD\", \"WY\"],\n",
    "    \"NE\": [\"CT\", \"ME\", \"MA\", \"NH\", \"RI\", \"VT\"],\n",
    "    \"SO\": [\"GA\", \"NC\", \"SC\", \"TN\", \"WV\", \"AR\"],\n",
    "    \"SW\": [\"AZ\", \"NV\", \"NM\", \"TX\", \"UT\"],\n",
    "    \"CN\": [\"AK\", \"IA\", \"KS\", \"MO\", \"NE\", \"OK\"],\n",
    "    \"WC\": [\"CA\", \"OR\", \"WA\"],\n",
    "}\n",
    "\n",
    "\n",
    "def map_region(state):\n",
    "    for key in regions:\n",
    "        if state in regions[key]:\n",
    "            return key\n",
    "\n",
    "\n",
    "data[\"REGION\"] = data.ST_ABB.map(map_region)\n",
    "data[\"TOTAL_EMP\"] = data.groupby([\"REGION\", \"YR\"])[\"EMP\"].transform(\"sum\")\n",
    "data[\"EMP_SHARE\"] = data.EMP / data.TOTAL_EMP\n",
    "data[\"WEIGHED_UNEMP\"] = data.EMP_SHARE * data.UNEMP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `groupby` transformation is used to aggregate the data, and finally all values except the unemployment rate are logged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grouped = data.groupby([\"REGION\", \"YR\"])\n",
    "agg_data = grouped[[\"GSP\", \"PC\", \"HWY\", \"WATER\", \"UTIL\", \"EMP\", \"WEIGHED_UNEMP\"]].sum()\n",
    "for col in [\"GSP\", \"PC\", \"HWY\", \"WATER\", \"UTIL\", \"EMP\"]:\n",
    "    agg_data[\"ln\" + col] = np.log(agg_data[col])\n",
    "agg_data[\"UNEMP\"] = agg_data.WEIGHED_UNEMP\n",
    "agg_data[\"Intercept\"] = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Usage\n",
    "\n",
    "Seemingly Unrelated Models are fairly complex and each equation could have a different number of regressors.  As a result, it is not possibly to use standard `pandas` or `numpy` data structures, and so dictionaries (or technically dictionary-like objects) are used.  In practice, it is strongly recommended to use a `OrderedDictionary` from the `collections` module.  This ensures that equation order will be preserved. In addition, the dictionary must have the following structure:\n",
    "\n",
    "* `keys` **must be strings** and will be used as equation labels\n",
    "* The value associated with each key must be either a dictionary or a tuple.\n",
    "\n",
    "  * When a dictionary is used, it must have two keys, `dependent` and `exog`.  It can optionally have a third key `weights` which provides weights to use in the regression.\n",
    "  * When a tuple is used, it must have two elements and takes the form `(dependent, exog)`.  It can optionally contains weights in which case it takes the form `(dependent, exog, weights)`.\n",
    "\n",
    "This example uses the dictionary syntax to contain the data for each region and uses the region identified as the equation label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "\n",
    "mod_data = OrderedDict()\n",
    "for region in [\"GF\", \"SW\", \"WC\", \"MT\", \"NE\", \"MA\", \"SO\", \"MW\", \"CN\"]:\n",
    "    region_data = agg_data.loc[region]\n",
    "    dependent = region_data.lnGSP\n",
    "    exog = region_data[\n",
    "        [\"Intercept\", \"lnPC\", \"lnHWY\", \"lnWATER\", \"lnUTIL\", \"lnEMP\", \"UNEMP\"]\n",
    "    ]\n",
    "    mod_data[region] = {\"dependent\": dependent, \"exog\": exog}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fitting the model is virtually identical to  fitting any other model with the exception of the special data structure required. \n",
    "\n",
    "The fitting options here ensure that the homoskedastic covariance estimator are used (`cov_type=\"unadjusted\"`) and that a small sample adjustment is applied. By default, GLS is used (this can be overridden using `method=\"ols\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from linearmodels.system import SUR\n",
    "\n",
    "mod = SUR(mod_data)\n",
    "res = mod.fit(cov_type=\"unadjusted\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the requirements for there to be an efficiency gain in a SUR is that the residuals are correlated. A heatmap is used to inspect this correlation, which is substantial and varies by region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "cov = res.sigma\n",
    "std = np.sqrt(np.diag(res.sigma)[:, None])\n",
    "regions = [k for k in mod_data.keys()]\n",
    "corr = pd.DataFrame(cov / (std @ std.T), columns=regions, index=regions)\n",
    "\n",
    "sns.heatmap(corr, vmax=0.8, square=True)\n",
    "plt.show()\n",
    "\n",
    "corr.style.format(\"{:0.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These values can be seen to be identical to the reported results in the existing example from Greene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display_png\n",
    "\n",
    "display_png(Image(\"system_correct-greene-table-10-2.png\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The full result is fairly long and so here I only print the first 33 lines which show results for two regions.  By default it reports all estimates along with the usual measures of precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\n\".join(res.summary.as_text().split(\"\\n\")[:33]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Individual results are contained in a dictionary located at the attribute `equations` and can be accessed using equation labels (available using the attribute `equation_labels`).  Additional information about the model is presented in this view. The West Coast results are show.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(res.equations[\"WC\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The current version of the model does not facilitate cross equation comparisons and so this is manually implemented here. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Implement method to compare across equations\n",
    "params = []\n",
    "for label in res.equation_labels:\n",
    "    params.append(res.equations[label].params)\n",
    "params = pd.concat(params, axis=1)\n",
    "params.columns = res.equation_labels\n",
    "params.T.style.format(\"{:0.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These results can be compared to the results in Greene -- they are unsurprisingly identical."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_png(Image(\"system_correct-greene-table-10-1.png\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The GLS estimation method requires stronger assumptions for parameter estimates to be consistent.  If these are violated then it might be the case that OLS is still consistent (in some sense) and so OLS can be used by passing `method=\"ols\"` when calling `fit`. \n",
    "\n",
    "These results can be compared to Greene\"s table -- they are identical except the final value which seems to have a small typo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_ols = mod.fit(method=\"ols\", debiased=True, cov_type=\"unadjusted\")\n",
    "params = []\n",
    "r2 = []\n",
    "for label in res.equation_labels:\n",
    "    params.append(res_ols.equations[label].params)\n",
    "    r2.append(res_ols.equations[label].rsquared)\n",
    "params = pd.concat(params, axis=1)\n",
    "params.columns = res.equation_labels\n",
    "params = params.T\n",
    "params[\"R2\"] = r2\n",
    "params.style.format(\"{:0.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The parameter estimates for one coefficient -- unemployment -- can be compared across the two estimation methods.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = pd.concat(\n",
    "    [\n",
    "        res_ols.params.iloc[1::7],\n",
    "        res_ols.std_errors.iloc[1::7],\n",
    "        res.params.iloc[1::7],\n",
    "        res.std_errors.iloc[1::7],\n",
    "    ],\n",
    "    axis=1,\n",
    ")\n",
    "params.columns = [\"OLS\", \"OLS se\", \"GLS\", \"GLS se\"]\n",
    "params.index = regions\n",
    "params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The parameters and their standard errors match those reported in Greene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_png(Image(\"system_correct-greene-table-10-3.png\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimation Options"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iterative GLS\n",
    "These next examples use data on fringe benefits from F. Vella (1993), \"A Simple Estimator for Simultaneous Models with Censored\n",
    "Endogenous Regressors\" which appears in Wooldridge (2002).  The model consists of two equations, one for hourly wage and the other for hourly benefits.  The initial model uses the same regressors in both equations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from linearmodels.datasets import fringe\n",
    "\n",
    "print(fringe.DESCR)\n",
    "fdata = fringe.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The estimated model is reported as \"System OLS Estimation\" since when all regressors are identical, OLS is used since GLS brings no efficiency gains.  OLS will be used when the data structure containing the exogenous regressors in each equations is the same (i.e., `id(exog)` is the same for all equations)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exog = sm.add_constant(\n",
    "    fdata[\n",
    "        [\n",
    "            \"educ\",\n",
    "            \"exper\",\n",
    "            \"expersq\",\n",
    "            \"tenure\",\n",
    "            \"tenuresq\",\n",
    "            \"union\",\n",
    "            \"south\",\n",
    "            \"nrtheast\",\n",
    "            \"nrthcen\",\n",
    "            \"married\",\n",
    "            \"white\",\n",
    "            \"male\",\n",
    "        ]\n",
    "    ]\n",
    ")\n",
    "fmod_data = OrderedDict()\n",
    "fmod_data[\"hrearn\"] = {\"dependent\": fdata.hrearn, \"exog\": exog}\n",
    "fmod_data[\"hrbens\"] = {\"dependent\": fdata.hrbens, \"exog\": exog}\n",
    "fmod = SUR(fmod_data)\n",
    "print(fmod.fit(cov_type=\"unadjusted\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The estimator can be forced to use GLS by setting `method=\"gls\"`.  As can be seen below, the parameter estimates and standard errors do not change even though two-stage GLS is used.  The $R^2$ **does change** since the left-hand side variable is transformed by the GLS weighting before estimation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(fmod.fit(method=\"gls\", cov_type=\"unadjusted\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The regressors must differ in order to see gains to GLS.  Here insignificant variables are dropped from each equation so that the regressors will no longer be identical. The typical standard error is marginally smaller, although this might be due to dropping regressors. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exog_earn = sm.add_constant(\n",
    "    fdata[[\"educ\", \"exper\", \"expersq\", \"union\", \"nrtheast\", \"white\"]]\n",
    ")\n",
    "exog_bens = sm.add_constant(\n",
    "    fdata[[\"educ\", \"exper\", \"expersq\", \"tenure\", \"tenuresq\", \"union\", \"male\"]]\n",
    ")\n",
    "fmod_data[\"hrearn\"] = {\"dependent\": fdata.hrearn, \"exog\": exog_earn}\n",
    "fmod_data[\"hrbens\"] = {\"dependent\": fdata.hrbens, \"exog\": exog_bens}\n",
    "fmod = SUR(fmod_data)\n",
    "print(fmod.fit(cov_type=\"unadjusted\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The standard method to estimate models uses two steps.  The first uses OLS to estimate the parameters so that the residual covariance can be estimated.  The second stage uses the estimated covariance to estimate the model via GLS.  Iterative GLS can be used to continue these iterations using the most recent step to estimate the residual covariance and then re-estimating the parameters using GLS.  This option can be used by setting `iterate=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fmod_res = fmod.fit(cov_type=\"unadjusted\", iterate=True)\n",
    "print(fmod_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of GLS iterations can be verified using the `iterations` attribute of the results. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fmod_res.iterations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alternative Covariance Estimators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The estimator supports heteroskedasticity robust covariance estimation b setting `cov_type=\"robust\"`.  The estimator allows for arbitrary correlation across series with the same time index but not correlation across time periods, which is the same assumption as in the unadjusted estimator. The main difference is that this estimator will correct standard errors for dependence between regressors (or squared regressors) and squared residuals.\n",
    "\n",
    "#### Heteroskedasticity Robust Covariance Estimation\n",
    "\n",
    "In the fringe benefit model there are some large differences between standard errors computed using the the homoskedastic covariance estimator and the heteroskedasticity robust covariance estimator (e.g., `exper`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fres_het = fmod.fit(cov_type=\"robust\")\n",
    "print(fres_het.summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Other supported covariance estimators include `\"kernel`\" which implements a HAC and `\"clustered\"` which supports 1 and 2-way clustering. \n",
    "\n",
    "#### Kernel (HAC)\n",
    "\n",
    "The supported kernels are `\"bartlett\"` (Newey-West), `\"parzen\"` (Gallant), and `qs` (Quadratic Spectral, Andrews). This example uses the Parzen kernel. The kernel's bandwidth is computed automatically if the parameter `bandwidth` is not provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hac_res = fmod.fit(cov_type=\"kernel\", kernel=\"parzen\")\n",
    "print(hac_res.summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Clustered (Rogers)\n",
    "\n",
    "The clustered covariance estimator requires the clusters to be entered as a NumPy array with shape  `(nobs, 1)` or `(nobs, 2)`. This example uses random clusters to illustrate the structure of the group id variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rs = np.random.RandomState([983476381, 28390328, 23829810])\n",
    "random_clusters = rs.randint(0, 51, size=(616, 1))\n",
    "clustered_res = fmod.fit(cov_type=\"clustered\", clusters=random_clusters)\n",
    "print(clustered_res.summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prespecified Residual Covariance Estimators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The GLS estimator can be used with a user specified covariance.  This example uses a covariance where all correlations are identical (equicorrelation) in the state GDP model.  The estimator must be used when constructing the model through the `sigma` keyword argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_corr = (corr - np.eye(9)).mean().mean() * (81 / 72)\n",
    "rho = np.ones((9, 9)) * avg_corr + (1 - avg_corr) * np.eye(9)\n",
    "sigma_pre = rho * (std @ std.T)\n",
    "\n",
    "mod_pre_sigma = SUR(mod_data, sigma=sigma_pre)\n",
    "res_pre = mod_pre_sigma.fit(cov_type=\"unadjusted\", debiased=True)\n",
    "print(res_pre.equations[\"GF\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Diagonal Residual Covariance\n",
    "\n",
    "The default assumption is that there can be arbitrary correlation between series.  The model can be estimated assuming no correlation but different variance by setting the keyword argument `full_cov=False` in the `fit` method.  This is only useful when there are cross-equation parameter restrictions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cross-Equation Restrictions\n",
    "\n",
    "One of the more useful features of using a SUR to estimate a system is the ability to impose constraints on parameters across equations.  Here only linear constraints of the form \n",
    "\n",
    "$$ r \\beta = q $$\n",
    "\n",
    "are supported.  Linear constraints are entered by passing a `DataFrame` with the shape number of constraints by  number of parameters.  The number and name of parameters can be seen by inspecting the `param_names` attribute of a model.  Below are the parameter names from the state GDP model which consist of the equation label, an underscore, and the variable name. This ensures uniqueness.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod.param_names[:14]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The example of a parameter constraint will enforce a common value of the coefficient of unemployment in all equations.  This restriction takes the form\n",
    "\n",
    "$$ \\beta_{unemp,i} - \\beta_{unemp,j} = 0$$\n",
    "\n",
    "where in all examples i is the first series and j is one of the 8 others.  In total there are 8 restrictions.  \n",
    "\n",
    "The construction of the restriction array and the non-zero columns are shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = pd.DataFrame(\n",
    "    columns=mod.param_names,\n",
    "    index=[\"rest{0}\".format(i) for i in range(8)],\n",
    "    dtype=np.float64,\n",
    ")\n",
    "r.loc[:, :] = 0.0\n",
    "r.iloc[:, 6] = -1.0\n",
    "r.iloc[:, 13::7] = np.eye(8)\n",
    "print(r.iloc[:, 6::7])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The constraint can be added to an existing model using the `add_constraint` method.  This method requires one input, `r` and optionally `q`.  if `q` is not provided, it is set to 0. \n",
    "\n",
    "Here the constraint is added, the model is estimated, and the parameters for unemployment are displayed.  They all have the same value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod.add_constraints(r)\n",
    "constrained_res = mod.fit()\n",
    "constrained_res.params[6::7]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pooling using constraints\n",
    "\n",
    "More complicated constraints can be used to produce interesting models.  Using the same idea as the previous set of constraints, a pooled SUR (excluding the constant) is constructed by restricting all coefficients to have the same value.  Here the form of the restriction is \n",
    "\n",
    "$$ \\beta_{var,0} - \\beta_{var,j} = 0$$\n",
    "\n",
    "so that the restriction is identical to the previous, only applied to all variables excluding the constant.\n",
    "\n",
    "Here the estimated from the first two equations are shown. All coefficients except the intercept are identical across equations.\n",
    "\n",
    "**Note**: When linear constraints are imposed, the parameter covariance matrix will be singular.  Caution is needed to ensure test statistics are meaningful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r2 = np.zeros((8 * 6, r.shape[1]))\n",
    "loc = 0\n",
    "for i in range(6):\n",
    "    for j in range(8):\n",
    "        r2[loc, i + 1] = -1.0\n",
    "        r2[loc, 7 * (j + 1) + i + 1] = 1.0\n",
    "        loc += 1\n",
    "r2 = pd.DataFrame(r2, columns=mod.param_names)\n",
    "mod.reset_constraints()\n",
    "mod.add_constraints(r2)\n",
    "pooled_res = mod.fit()\n",
    "print(\"\\n\".join(pooled_res.summary.as_text().split(\"\\n\")[:33]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multivariate OLS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One final feature worth demonstrating is a simple interface for specifying multivariate OLS models.  These models have the same regressors and so it is possible to specify them with two arrays. The first is a $T$ by $K$ array of dependent variables where each column contains a dependent variable.  The second contains the common exogenous regressors.\n",
    "\n",
    "This example shows how a CAPM can be estimated as a MV OLS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import statsmodels.api as sm\n",
    "\n",
    "from linearmodels.datasets import french\n",
    "\n",
    "data = french.load()\n",
    "factors = sm.add_constant(data[[\"MktRF\"]])\n",
    "mv_ols = SUR.multivariate_ls(\n",
    "    data[[\"S1V1\", \"S1V3\", \"S1V5\", \"S5V1\", \"S5V3\", \"S5V5\"]], factors\n",
    ")\n",
    "mv_ols_res = mv_ols.fit(cov_type=\"unadjusted\")\n",
    "print(mv_ols_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using GLS with common regressors\n",
    "\n",
    "As noted previously, forcing GLS has no effect (except on changing the $R^2$).  This can be seen below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mv_ols.fit(cov_type=\"unadjusted\", method=\"gls\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
